/**
 * 2698. 求一个整数的惩罚数
 */
public class No2698 {
    public int punishmentNumber(int n) {
        int ans = 0;
        for (int i = 0; i <= n; i++) {
            if (isPunishment(i, (i * i + "").length(), 0)) {
                ans += i * i;
            }
        }
        return ans;
    }

    private boolean isPunishment(int i, int endIndex, int curSum) {
        if (curSum > i) return false;
        else if (endIndex == 0) return curSum == i;
        String multi = i * i + "";
        for (int startIndex = endIndex - 1; startIndex >= 0; startIndex--) {
            int curNum = Integer.parseInt(multi.substring(startIndex, endIndex));
            curSum += curNum;
            if (isPunishment(i, startIndex, curSum)) return true;
            curSum -= curNum;
        }
        return false;
    }
}
